import numpy as np
from scipy.spatial.transform import Rotation


# 向量通过矩阵变换
def transform3(m, v):
    return transform_vec3(m, v) + m[..., :3, 3]


def transform_vec3(m, v):
    if m.ndim == 2:
        return (v.reshape(-1, 3) @ m[:3, :3].T).reshape(v.shape)
    else:
        return (m[..., :3, :3] @ v[..., None]).squeeze(-1)


# 正则化
def normalized(v: np.ndarray, axis: int = -1, eps: float = 5.43e-20) -> np.ndarray:
    d = np.maximum(eps, (v * v).sum(axis=axis, keepdims=True) ** 0.5)
    return v / d


# 生成向量的叉乘矩阵
def skew_matrix(v: np.ndarray) -> np.ndarray:
    res = np.array(
        [[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]], dtype=v.dtype
    )
    return res


# 通过 Rodrigues 公式生成一个最小旋转矩阵，把 a 旋转到 b
def from_two_vectors(a_orig: np.ndarray, b_orig: np.ndarray) -> np.ndarray:
    a = normalized(a_orig)
    b = normalized(b_orig)
    v = np.cross(a, b)
    s = np.linalg.norm(v)
    c = np.dot(a, b)
    v_mat = skew_matrix(v)

    rot = np.eye(3, 3) + v_mat + np.matmul(v_mat, v_mat) * (1 - c) / (max(s * s, 1e-15))

    return rot


# 生成lookat矩阵，用来将世界坐标系转为相机坐标系
import numpy as np
from scipy.spatial.transform import Rotation


def make_look_at_matrix(
        orig_world_to_eye: np.ndarray,
        center: np.ndarray,
        camera_angle: float = 0,
        augment: bool = False,
        angle_perturb_range: float = 15.0,
        dir_perturb_range: float = 0.1
) -> np.ndarray:
    # 将中心点转换到相机坐标系
    center_local = transform3(orig_world_to_eye, center.copy())

    # 计算指向中心的方向向量并归一化
    z_dir_local = center_local / np.linalg.norm(center_local)

    # 数据增强：在归一化方向向量上添加扰动
    if augment:
        # 创建随机扰动向量（在相机坐标系XY平面）
        perturb = np.zeros(3, dtype=center.dtype)
        perturb[0] = np.random.uniform(-dir_perturb_range, dir_perturb_range)
        perturb[1] = np.random.uniform(-dir_perturb_range, dir_perturb_range)

        # 应用扰动并重新归一化
        z_dir_local += perturb
        z_dir_local /= np.linalg.norm(z_dir_local)

    # 计算旋转矩阵使相机Z轴指向扰动后的方向
    delta_r_local = from_two_vectors(
        np.array([0, 0, 1], dtype=center.dtype), z_dir_local
    )

    # 数据增强：相机角度扰动
    if augment:
        camera_angle += np.random.uniform(-angle_perturb_range, angle_perturb_range)

    # 获取原始相机->世界的变换矩阵
    orig_eye_to_world = np.linalg.inv(orig_world_to_eye)

    # 创建新的相机->世界变换矩阵
    new_eye_to_world = orig_eye_to_world.copy()
    new_eye_to_world[0:3, 0:3] = orig_eye_to_world[0:3, 0:3] @ delta_r_local

    # 添加绕Z轴的旋转
    z_local_rot = Rotation.from_euler("z", camera_angle, degrees=True).as_matrix()
    new_eye_to_world[0:3, 0:3] = new_eye_to_world[0:3, 0:3] @ z_local_rot

    # 返回新的世界->相机变换矩阵
    return np.linalg.inv(new_eye_to_world)



def make_look_at_from_camera_space(
        orig_world_to_eye: np.ndarray,
        local_target: np.ndarray,  # 目标点在相机坐标系中的位置
        camera_angle: float = 0,
) -> np.ndarray:
    """
    在保持相机位置不变的情况下，旋转相机使其对准相机坐标系中的目标点

    参数:
        orig_world_to_eye: 原始的世界坐标系→相机坐标系的4x4变换矩阵
        local_target: 目标点在相机坐标系中的位置 (3D向量)
        camera_angle: 相机绕自身光轴(Z轴)的旋转角度(度)

    返回:
        新的世界坐标系→相机坐标系的4x4变换矩阵
    """
    # 归一化目标方向向量 (新Z轴方向)
    z_dir_local = local_target / np.linalg.norm(local_target)

    # 计算对准旋转矩阵 (从原始Z轴到新Z轴)
    delta_r_local = from_two_vectors(
        np.array([0, 0, 1], dtype=local_target.dtype),
        z_dir_local
    )

    # 获取原始相机→世界变换矩阵
    orig_eye_to_world = np.linalg.pinv(orig_world_to_eye)

    # 创建新的相机→世界变换矩阵
    new_eye_to_world = orig_eye_to_world.copy()

    # 应用对准旋转
    new_eye_to_world[0:3, 0:3] = orig_eye_to_world[0:3, 0:3] @ delta_r_local

    # 应用光轴旋转 (如果需要)
    if camera_angle != 0:
        z_local_rot = Rotation.from_euler("z", camera_angle, degrees=True).as_matrix()
        new_eye_to_world[0:3, 0:3] = new_eye_to_world[0:3, 0:3] @ z_local_rot

    # 返回新的世界→相机变换矩阵
    return np.linalg.pinv(new_eye_to_world)
